{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d006b2ea-9dfe-49c7-88a9-a5a0775185fd",
   "metadata": {},
   "source": [
    "# Additional End of week Exercise - week 2\n",
    "\n",
    "Now use everything you've learned from Week 2 to build a full prototype for the technical question/answerer you built in Week 1 Exercise.\n",
    "\n",
    "This should include a Gradio UI, streaming, use of the system prompt to add expertise, and the ability to switch between models. Bonus points if you can demonstrate use of a tool!\n",
    "\n",
    "If you feel bold, see if you can add audio input so you can talk to it, and have it respond with audio. ChatGPT or Claude can help you, or email me if you have questions.\n",
    "\n",
    "I will publish a full solution here soon - unless someone beats me to it...\n",
    "\n",
    "There are so many commercial applications for this, from a language tutor, to a company onboarding solution, to a companion AI to a course (like this one!) I can't wait to see your results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1989a03e-ed40-4b8c-bddd-322032ca99f5",
   "metadata": {},
   "source": [
    "# Advanced Airline AI Assistant\n",
    "### original features:\n",
    "1. chat with the AI assistant\n",
    "2. use a Tool to get ticket price\n",
    "3. generate Audio for each AI response \n",
    "### advanced features:\n",
    "3. add a Tool to make a booking\n",
    "4. add an Agent that translate all responses to a different language\n",
    "5. add an Agent that can listen for Audio and convert to Text\n",
    "6. generate audio for each user input and AI response, including both the original and translated versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ed79822-af6b-4bfb-b108-5f36e237e97a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Library for language translation\n",
    " \n",
    "!pip install deep_translator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29184b81-b945-4dd3-bd17-2c64466d37d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Library for speech-to-text conversion\n",
    "# make sure 'ffmpeg' is downloaded already\n",
    "\n",
    "!pip install openai-whisper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2b0a9b2-ce83-42ff-a312-582dc5ee9097",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Library for storing and loading audio file\n",
    "\n",
    "pip install soundfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07e7793-b8f5-44f4-aded-5562f633271a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import gradio as gr\n",
    "import base64\n",
    "from io import BytesIO\n",
    "from IPython.display import Audio, display\n",
    "import tempfile\n",
    "import whisper\n",
    "import soundfile as sf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da46ca14-2052-4321-a940-2f2e07b40975",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialization\n",
    "\n",
    "load_dotenv(override=True)\n",
    "\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "if openai_api_key:\n",
    "    print(f\"OpenAI API Key exists and begins {openai_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"OpenAI API Key not set\")\n",
    "    \n",
    "MODEL = \"gpt-4o-mini\"\n",
    "openai = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "499d3d06-9628-4a69-bc9d-fa481fd8fa98",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"You are a helpful assistant for an Airline called FlightAI. \"\n",
    "system_message += \"Your main responsibilities are solve customers' doubts, get ticket price and book a ticket\"\n",
    "system_message += \"Give short, courteous answers, no more than 1 sentence. \"\n",
    "system_message += \"Always be accurate. If you don't know the answer, say so.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25cf964e-a954-43d5-85bd-964efe502c25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's start by making a useful function\n",
    "\n",
    "ticket_prices = {\"london\": \"$799\", \"paris\": \"$899\", \"tokyo\": \"$1400\", \"berlin\": \"$499\", \"shanghai\": \"$799\", \"wuhan\": \"$899\"}\n",
    "\n",
    "def get_ticket_price(destination_city):\n",
    "    print(f\"Tool get_ticket_price called for {destination_city}\")\n",
    "    city = destination_city.lower()\n",
    "    return ticket_prices.get(city, \"Unknown\")\n",
    "\n",
    "def book_ticket(destination_city):\n",
    "    print(f\"Tool book_ticket called for {destination_city}\")\n",
    "    city = destination_city.lower()\n",
    "    global booked_cities\n",
    "    if city in ticket_prices:\n",
    "        price = ticket_prices.get(city, \"\")\n",
    "        label = f\"{city.title()} ({price})\"\n",
    "        i = booked_cities_choices.index(city.lower().capitalize())\n",
    "        booked_cities_choices[i] = label\n",
    "        booked_cities.append(label)\n",
    "        return f\"Booking confirmed for {city.title()} at {ticket_prices[city]}\"\n",
    "    else:\n",
    "        return \"City not found in ticket prices.\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "701aa037-1ab3-4861-a809-b7f13ef9ea36",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# There's a particular dictionary structure that's required to describe our function:\n",
    "\n",
    "price_function = {\n",
    "    \"name\": \"get_ticket_price\",\n",
    "    \"description\": \"Get the price of a return ticket to the destination city. Call this whenever you need to know the ticket price, for example when a customer asks 'How much is a ticket to this city'\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"destination_city\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The city that the customer wants to travel to\",\n",
    "            },\n",
    "        },\n",
    "        \"required\": [\"destination_city\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}\n",
    "\n",
    "book_function = {\n",
    "    \"name\": \"book_ticket\",\n",
    "    \"description\": \"Book a return ticket to the destination city. Call this whenever you want to book a ticket to the city, for example when the user says something like 'Book me a ticket to this city'\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"destination_city\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The city that the customer wants to book a ticket to\"\n",
    "            }\n",
    "        },\n",
    "        \"required\": [\"destination_city\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c4cf01c-ba15-4a4b-98db-6f86c712ec66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And this is included in a list of tools:\n",
    "\n",
    "tools = [\n",
    "    {\"type\": \"function\", \"function\": price_function},\n",
    "    {\"type\": \"function\", \"function\": book_function}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7486e2c-4687-4819-948d-487b5e528fc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydub import AudioSegment\n",
    "from pydub.playback import play\n",
    "\n",
    "def talker(message):\n",
    "    response = openai.audio.speech.create(\n",
    "      model=\"tts-1\",\n",
    "      voice=\"onyx\",    # Also, try replacing onyx with alloy\n",
    "      input=message\n",
    "    )\n",
    "    \n",
    "    audio_stream = BytesIO(response.content)\n",
    "    audio = AudioSegment.from_file(audio_stream, format=\"mp3\")\n",
    "    play(audio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac195914-4a89-462c-9be0-fee286498491",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This part is inspired from 'week2/community-contributions/week2_exerccise_translated_chatbot'\n",
    "from deep_translator import GoogleTranslator\n",
    "\n",
    "# Available translation language\n",
    "LANGUAGES = {\n",
    "    \"English\": \"en\",\n",
    "    \"Mandarin Chinese\": \"zh-CN\",\n",
    "    \"Hindi\": \"hi\",\n",
    "    \"Spanish\": \"es\",\n",
    "    \"Arabic\": \"ar\",\n",
    "    \"Bengali\": \"bn\",\n",
    "    \"Portuguese\": \"pt\",\n",
    "    \"Russian\": \"ru\",\n",
    "    \"Japanese\": \"ja\",\n",
    "    \"German\": \"de\"\n",
    "}\n",
    "\n",
    "def update_lang(choice):\n",
    "    global target_lang\n",
    "    target_lang = LANGUAGES.get(choice, \"zh-CN\") \n",
    "\n",
    "def translate_message(text, target_lang):\n",
    "    if target_lang == \"en\":\n",
    "        return text\n",
    "    try:\n",
    "        translator = GoogleTranslator(source='auto', target=target_lang)\n",
    "        return translator.translate(text)\n",
    "    except:\n",
    "        return f\"Translation error: {text}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46255fe5-9621-47ba-af78-d0c74aee2997",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Text-to-speech conversion\n",
    "def speak(message):\n",
    "    response = openai.audio.speech.create(\n",
    "        model=\"tts-1\",\n",
    "        voice=\"onyx\",\n",
    "        input=message)\n",
    "\n",
    "    audio_stream = BytesIO(response.content)\n",
    "    output_filename = \"output_audio.mp3\"\n",
    "    with open(output_filename, \"wb\") as f:\n",
    "        f.write(audio_stream.read())\n",
    "\n",
    "    # Play the generated audio\n",
    "    display(Audio(output_filename, autoplay=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d73f0b3a-34ae-4685-8a5d-8b6421f872c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update dropdown options from chatbot history\n",
    "def update_options(history):\n",
    "    options = [f\"{msg['role']}: {msg['content']}\" for msg in history]\n",
    "    return gr.update(choices=options, value=options[-1] if options else \"\")\n",
    "\n",
    "# Extract just the text content from selected entry\n",
    "def extract_text(selected_option):\n",
    "    return selected_option.split(\": \", 1)[1] if \": \" in selected_option else selected_option"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab12d51b-c799-4ce4-87d5-9ae2265d148f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Handles audio input as numpy array and returns updated chat history\n",
    "def speak_send(audio_np, history):\n",
    "    if audio_np is None:\n",
    "        return history\n",
    "\n",
    "    # Convert NumPy audio to in-memory .wav file\n",
    "    sample_rate, audio_array = audio_np\n",
    "    with tempfile.NamedTemporaryFile(suffix=\".wav\") as f:\n",
    "        sf.write(f.name, audio_array, sample_rate)\n",
    "        result = model.transcribe(f.name)\n",
    "        text = result[\"text\"]\n",
    "        \n",
    "    history += [{\"role\":\"user\", \"content\":text}]\n",
    "\n",
    "    return None, history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "221b1380-c894-45d4-aad2-e94b3b9454b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We have to write that function handle_tool_call:\n",
    "\n",
    "def handle_tool_call(message):\n",
    "    tool_call = message.tool_calls[0]\n",
    "    tool_name = tool_call.function.name\n",
    "    arguments = json.loads(tool_call.function.arguments)\n",
    "\n",
    "    if tool_name == \"get_ticket_price\":\n",
    "        city = arguments.get(\"destination_city\")\n",
    "        price = get_ticket_price(city)\n",
    "        response = {\n",
    "            \"role\": \"tool\",\n",
    "            \"content\": json.dumps({\"destination_city\": city,\"price\": price}),\n",
    "            \"tool_call_id\": tool_call.id\n",
    "        }\n",
    "        return response, city\n",
    "\n",
    "    elif tool_name == \"book_ticket\":\n",
    "        city = arguments.get(\"destination_city\")\n",
    "        result = book_ticket(city)\n",
    "        response = {\n",
    "            \"role\": \"tool\",\n",
    "            \"content\": result,\n",
    "            \"tool_call_id\": tool_call.id            \n",
    "        }\n",
    "        return response, city\n",
    "\n",
    "    else:\n",
    "        return {\n",
    "            \"role\": \"tool\",\n",
    "            \"content\": f\"No tool handler for {tool_name}\",\n",
    "            \"tool_call_id\": tool_call.id\n",
    "        }, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27f19cd3-53cd-4da2-8be0-1fdd5424a7c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The advanced 'chat' function in 'day5'\n",
    "def interact(history, translated_history):\n",
    "    messages = [{\"role\": \"system\", \"content\": system_message}] + history\n",
    "    response = openai.chat.completions.create(model=MODEL, messages=messages, tools=tools)\n",
    "    \n",
    "    if response.choices[0].finish_reason==\"tool_calls\":\n",
    "        message = response.choices[0].message\n",
    "        response, city = handle_tool_call(message)\n",
    "        messages.append(message)\n",
    "        messages.append(response)\n",
    "        response = openai.chat.completions.create(model=MODEL, messages=messages)\n",
    "        \n",
    "    reply = response.choices[0].message.content\n",
    "    translated_message = translate_message(history[-1][\"content\"], target_lang)\n",
    "    translated_reply = translate_message(reply, target_lang)\n",
    "    \n",
    "    history += [{\"role\":\"assistant\", \"content\":reply}]\n",
    "    translated_history += [{\"role\":\"user\", \"content\":translated_message}]\n",
    "    translated_history += [{\"role\":\"assistant\", \"content\":translated_reply}]\n",
    "    \n",
    "    # Comment out or delete the next line if you'd rather skip Audio for now..\n",
    "    talker(reply)\n",
    "\n",
    "    return history, update_options(history), history, translated_history, update_options(translated_history), translated_history, gr.update(choices=booked_cities_choices, value=booked_cities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f714b955-4fb5-47df-805b-79f813f97548",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks() as demo:\n",
    "    target_lang = \"zh-CN\"\n",
    "    history_state = gr.State([]) \n",
    "    translated_history_state = gr.State([])\n",
    "    booked_cities_choices = [key.lower().capitalize() for key in ticket_prices.keys()]\n",
    "    booked_cities = []\n",
    "    model = whisper.load_model(\"base\")\n",
    "\n",
    "    with gr.Row():\n",
    "        city_checklist = gr.CheckboxGroup(\n",
    "            label=\"Booked Cities\",\n",
    "            choices=booked_cities_choices     \n",
    "        )\n",
    "            \n",
    "    with gr.Row():\n",
    "        with gr.Column():\n",
    "            chatbot = gr.Chatbot(label=\"Chat History\", type=\"messages\")\n",
    "            selected_msg = gr.Dropdown(label=\"Select message to speak\", choices=[])\n",
    "            speak_btn = gr.Button(\"Speak\")\n",
    "\n",
    "        with gr.Column():\n",
    "            translated_chatbot = gr.Chatbot(label=\"Translated Chat History\", type=\"messages\")\n",
    "            translated_selected_msg = gr.Dropdown(label=\"Select message to speak\", choices=[], interactive=True)\n",
    "            translated_speak_btn = gr.Button(\"Speak\")\n",
    "    \n",
    "    with gr.Row():\n",
    "        language_dropdown = gr.Dropdown(\n",
    "                choices=list(LANGUAGES.keys()),\n",
    "                value=\"Mandarin Chinese\",\n",
    "                label=\"Translation Language\",\n",
    "                interactive=True\n",
    "            )\n",
    "      \n",
    "    with gr.Row():\n",
    "        entry = gr.Textbox(label=\"Chat with our AI Assistant:\")\n",
    "\n",
    "    with gr.Row():\n",
    "        audio_input = gr.Audio(sources=\"microphone\", type=\"numpy\", label=\"Speak with our AI Assistant:\")\n",
    "    with gr.Row():\n",
    "        audio_submit = gr.Button(\"Send\")\n",
    "    \n",
    "    def do_entry(message, history):\n",
    "        history += [{\"role\":\"user\", \"content\":message}]\n",
    "        return \"\", history\n",
    "        \n",
    "    language_dropdown.change(fn=update_lang, inputs=[language_dropdown])\n",
    "\n",
    "    speak_btn.click(\n",
    "        lambda selected: speak(extract_text(selected)),\n",
    "        inputs=selected_msg,\n",
    "        outputs=None\n",
    "    )\n",
    "\n",
    "    translated_speak_btn.click(\n",
    "        lambda selected: speak(extract_text(selected)),\n",
    "        inputs=translated_selected_msg,\n",
    "        outputs=None\n",
    "    )\n",
    "\n",
    "    entry.submit(do_entry, inputs=[entry, history_state], outputs=[entry, chatbot]).then(\n",
    "        interact, inputs=[chatbot, translated_chatbot], outputs=[chatbot, selected_msg, history_state, translated_chatbot, translated_selected_msg, translated_history_state, city_checklist]\n",
    "    )\n",
    "    \n",
    "    audio_submit.click(speak_send, inputs=[audio_input, history_state], outputs=[audio_input, chatbot]).then(\n",
    "        interact, inputs=[chatbot, translated_chatbot], outputs=[chatbot, selected_msg, history_state, translated_chatbot, translated_selected_msg, translated_history_state, city_checklist]\n",
    "    )\n",
    "    # clear.click(lambda: None, inputs=None, outputs=chatbot, queue=False)\n",
    "\n",
    "demo.launch()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
